import numpy as np
import my_assert as my
import random as rand
import time
import math as m

def ufunc(d=(100,100,100),c=3):
	ct = 0
	ct_nocopy = 0
	nt = 0
	for i in np.arange(c):
		ca = np.zeros(d, dtype=np.double, cuda=1)
		cb = np.zeros(d, dtype=np.double, cuda=1)
		na = np.zeros(d, dtype=np.double, cuda=0)
		nb = np.zeros(d, dtype=np.double, cuda=0)
		np.ufunc_random(ca,ca)
		np.ufunc_random(cb,cb)
		np.ufunc_random(na,na)
		np.ufunc_random(nb,nb)
		#Including device transfer
		t0=time.time()
		np.divide(ca,cb)
		t1=time.time()
		np.divide(na,nb)
		t2=time.time()
		#print 'hertil'
		ct += t1-t0
		nt += t2-t1
		#Excluding device transfer
		t0=time.time()
		np.divide(ca,cb)
		t1=time.time()
		ct_nocopy += t1-t0
	print 'ct:',ct/c,'ct_nocopy:',ct_nocopy/c,'nt:',nt/c


"""
print 'ufunc((i,100,100), c=6)'
for i in np.arange(100, 1001, 100):
	ufunc((i,10,100), c=6)
print 'ufunc((100,i,100), c=6)'
for i in np.arange(100, 1001, 100):
	ufunc((100,i,100), c=6)
print 'ufunc((100,100,i), c=6)' 
for i in np.arange(100, 1001, 100):
	ufunc((100,100,i), c=6)
"""

def fft_2d(x,y,c=3):
	ctR2C = 0
	ctC2C = 0
	ctC2R = 0
	ntR2C = 0
	ntC2C = 0
	ntC2R = 0
	for i in np.arange(c):
		cuda_A  = np.zeros((x,y), dtype=np.double, cuda=1)
		np.ufunc_random(cuda_A, cuda_A)
		numpy_A = my.copy_2d(cuda_A)
		np.add(cuda_A,1)#Force cuda_A onto the device
		#real -> complx
		t0 = time.time()
		cuda_B = np.cufftR2C(cuda_A)
		t1 = time.time()
		numpy_B = np.fft.rfftn(numpy_A)
		t2 = time.time()
		ctR2C += t1-t0
		ntR2C += t2-t1	

		#complex -> complex
		t0 = time.time()
		cuda_C = np.cufftC2C(cuda_B)
		t1 = time.time()
		numpy_C = np.fft.fftn(numpy_B)
		t2 = time.time()
		ctC2C += t1-t0
		ntC2C += t2-t1

		#complex -> real
		t0 = time.time()
		cuda_C = np.cufftC2R(cuda_B)/(x*y)
		t1 = time.time()
		numpy_C = np.fft.irfftn(numpy_B)
		t2 = time.time()
		ctC2R += t1-t0
		ntC2R += t2-t1

	print 'R2C cuda:',ctR2C/c,'numpy:',ntR2C/c,'C2C cuda:',ctC2C/c,'numpy:',ntC2C/c,'C2R cuda:',ctC2R/c,'numpy:',ntC2R/c
	return 1


def fft_1d(x,c=3):
	ctR2C = 0
	ctC2C = 0
	ctC2R = 0
	ntR2C = 0
	ntC2C = 0
	ntC2R = 0
	for i in np.arange(c):
		cuda_A = np.zeros((x), dtype=np.double, cuda=1)
		np.ufunc_random(cuda_A, cuda_A)
		numpy_A = my.copy_1d(cuda_A)
		np.add(cuda_A,1)#Force cuda_A onto the device
		#real -> complex
		t0 = time.time()
		cuda_res = np.cufftR2C(cuda_A)
		t1 = time.time()
		numpy_res = np.fft.fft(numpy_A)
		t2 = time.time()
		ctR2C += t1-t0
		ntR2C += t2-t1
		#print 'real -> complex cuda:',t1-t0,'numpy:',t2-t1

		#complex -> complex
		t0 = time.time()
		ct = np.cufftC2C(cuda_res)
		t1 = time.time()
		nt = np.fft.fft(numpy_res)
		t2 = time.time()
		ctC2C += t1-t0
		ntC2C += t2-t1
		#print 'complex -> complex cuda:',t1-t0,'numpy:',t2-t1

		#complex -> real
		t0 = time.time()
		ct = (np.cufftC2R(cuda_res))/x
		t1 = time.time()
		nt = np.fft.ifft(numpy_res)
		t2 = time.time()
		ctC2R += t1-t0
		ntC2R += t2-t1
		#print 'complex -> real cuda:',t1-t0,'numpy:',t2-t1

	print 'R2C cuda:',ctR2C/c,'numpy:',ntR2C/c,'C2C cuda:',ctC2C/c,'numpy:',ntC2C/c,'C2R cuda:',ctC2R/c,'numpy:',ntC2R/c
#fft_1d(500000)
#fft_1d(524288)
#fft_1d(600000)
#for i in np.arange(100000,4000000,200000):
#	fft_1d(i)
#for i in np.arange(100,2000,200):
#	fft_2d(i,i)
#fft_1d(1048576)
#fft_1d(1050000)

def dot_3d(m,k,n,c=3):
	ct=0
	nt=0
	for i in np.arange(c):
		cuda_A = np.zeros((m, k), dtype=np.double, cuda=1)
		cuda_B = np.zeros((k, n), dtype=np.double, cuda=1)
		np.ufunc_random(cuda_A, cuda_A)
		np.ufunc_random(cuda_B, cuda_B)
		np.add(cuda_A, cuda_B) #Force a & b on the device
		numpy_A = my.copy_2d(cuda_A)
		numpy_B = my.copy_2d(cuda_B)

		t0 = time.time()
		cuda_res=np.dot(cuda_A, cuda_B)
		t1 = time.time()
		numpy_res=np.dot(numpy_A, numpy_B)
		t2 = time.time()
		ct+=t1-t0
		nt+=t2-t1
	print m*k,'cuda:',ct/c,'numpy',nt/c
#dot_3d(1414,1414,1414)
#for i in np.arange(100,1100,100):
#	dot_3d(i,i,i)

def reduce_1d(x=100, c=3):
	nt=0
	ct=0
	for i in np.arange(c):
		cuda_A = np.zeros((x), dtype=np.double, cuda=1)
		numpy_A = np.zeros((x), dtype=np.double)
		np.ufunc_random(cuda_A, cuda_A)
		np.ufunc_random(numpy_A, numpy_A)
		#numpy_A+1
		t0=time.time()
		np.add.reduce(cuda_A)
		t1=time.time()
		np.add.reduce(numpy_A)
		t2=time.time()
		ct += t1-t0
		nt += t2-t1
	print '1D cuda',ct/c,'numpy',nt/c

def reduce_2d(x=100,y=100, c=3, a=0):
	nt=0
	ct=0
	for i in np.arange(c):
		cuda_A = np.zeros((x,y), dtype=np.double, cuda=1)
		numpy_A = np.zeros((x,y), dtype=np.double)
		np.ufunc_random(cuda_A, cuda_A)
		np.ufunc_random(numpy_A, numpy_A)
		#cuda_A+1
		#numpy_A+1
		t0=time.time()
		np.add.reduce(cuda_A, axis=a)
		t1=time.time()
		np.add.reduce(numpy_A, axis=a)
		t2=time.time()
		ct += t1-t0
		nt += t2-t1
	print '2D size:',x,'x',y,' cuda',ct/c,'numpy',nt/c


#reduce_1d(5121)
#for i in np.arange(10000,2200000,200000):
#	reduce_1d(i)


for i in np.arange(100,3101,500):
	#print (i*i)/1000
	reduce_2d(i,i,c=6)
for i in np.arange(100,18200,1800):
	#print (i*500)/1000
	reduce_2d(i,500,c=6)
for i in np.arange(100,18200,1800):
	#print i*500
	reduce_2d(500,i,c=6)


#for i in np.arange(500,7000,700):
#	reduce_2d(i,i)

#print '2000x200'
#reduce_2d(2000,200, a=0, c=6)
#reduce_2d(2000,200, a=1, c=6)
#print '2000x2000'
#reduce_2d(2000,2000, a=0, c=6)
#reduce_2d(2000,2000, a=1, c=6)

